import logging

from fastapi import APIRouter, Body, Query
from tortoise.expressions import Q

from app.controllers.datasets import dataset_controller
from app.schemas.base import Fail, Success, SuccessExtra
from app.schemas.models import *

logger = logging.getLogger(__name__)

router = APIRouter()


@router.get("/list", summary="查看数据集列表")
async def list_model(
    page: int = Query(1, description="页码"),
    page_size: int = Query(10, description="每页数量"),
    name: str = Query("", description="名称，用于搜索"),
    desc: str = Query("", description="数据集描述"),
    file_type: str = Query("", description="文件类型 是属于image or video or audio"),
):
    q = Q()
    if name:
        q &= Q(name__contains=name)
    if desc:
        q &= Q(desc__contains=desc)
    if file_type:
        q &= Q(file_type=file_type)
    total, model_objs = await dataset_controller.list(page=page, page_size=page_size, search=q)
    data = [await obj.to_dict(m2m=True, exclude_fields=[]) for obj in model_objs]
    return SuccessExtra(data=data, total=total, page=page, page_size=page_size)


@router.get("/get", summary="查看数据集")
async def get_model(
    model_id: int = Query(..., description="数据集ID"),
):
    model_obj = await dataset_controller.get(id=model_id)
    model_dict = await model_obj.to_dict(exclude_fields=["tasks","files"])
    return Success(data=model_dict)


@router.post("/create", summary="创建数据集")
async def create_model(
    model_in: ModelCreate,
):
    model = await dataset_controller.get_by_name(model_in.name)
    if model:
        return Fail(code=400, msg="系统中已存在带有此名称的数据集。")
    new_model = await dataset_controller.create_model(obj_in=model_in)
    model_dict = await new_model.to_dict(exclude_fields=[])
    return Success(msg="Created Successfully",data=model_dict)


@router.post("/update", summary="更新数据集")
async def update_model(
    model_in: ModelUpdate,
):
    model = await dataset_controller.update(id=model_in.id, obj_in=model_in)
    model_dict = await model.to_dict(exclude_fields=[])
    return Success(msg="Updated Successfully",data=model_dict)


@router.delete("/delete", summary="删除数据集")
async def delete_model(
    model_id: int = Query(..., description="数据集ID"),
):
    await dataset_controller.remove(id=model_id)
    return Success(msg="Deleted Successfully")

